import re
import pandas as pd

input_path = "../data/Titanic/test.csv"
df = pd.read_csv(input_path)

data = df.copy()
# 众数
embarked_mode = data["Embarked"].mode()[0]
data["Embarked"] = data["Embarked"].fillna(embarked_mode)

# 检查"data"数据框中"Cabin"列的每个值是否非空（notna()）
# 将布尔值结果转换为整数类型（astype(int)），True转为1，False转为0
# 将结果存储在新创建的"HasCabin"列中
data["HasCabin"] = data["Cabin"].notna().astype(int)

# 中位数
fare_median = data["Fare"].median()
data["Fare"] = data["Fare"].fillna(fare_median)


def extract_title(name: str) -> str:
    m = re.search(r",\s*([^\.]+)\.", name)
    return m.group(1).strip() if m else "Unknown"


data["Title"] = data["Name"].apply(extract_title)
title_counts = data["Title"].value_counts()
rare_titles = title_counts[title_counts < 10].index
data["Title"] = data["Title"].replace(rare_titles, "Rare")

age_median_by_title = data.groupby("Title")["Age"].median()


def impute_age(row):
    if pd.isna(row["Age"]):
        return age_median_by_title.get(row["Title"], data["Age"].median())
    return row["Age"]


data["Age"] = data.apply(impute_age, axis=1)

data["FamilySize"] = data["SibSp"] + data["Parch"] + 1

data["IsAlone"] = (data["FamilySize"] == 1).astype(int)

data["Sex_bin"] = data["Sex"].map({"male": 0, "female": 1}).astype(int)

embarked_ohe = pd.get_dummies(data["Embarked"], prefix="Embarked", drop_first=False)
pclass_ohe = pd.get_dummies(data["Pclass"].astype("category"), prefix="Pclass", drop_first=False)

data = pd.concat([data, embarked_ohe, pclass_ohe], axis=1)

data["Fare_z"] = (data["Fare"] - data["Fare"].mean()) / data["Fare"].std(ddof=0)
data["Age_z"] = (data["Age"] - data["Age"].mean()) / data["Age"].std(ddof=0)

data = data.drop(columns=["Name", "Ticket", "Cabin"])

# front_cols = ["PassengerId", "Survived"]
front_cols = ["PassengerId"]
other_cols = [c for c in data.columns if c not in front_cols]
data = data[front_cols + other_cols]

print("各列缺失值数量：")
print(data.isna().sum())

print("\n处理后数据预览（前 5 行）：")
print(data.head())

print("\n处理后数据形状：", data.shape)

data.to_csv("../data/Titanic/titanic_test1.csv", index=False)
